/*
** Copyright (C) 2010 Calin Radoni
**
** This file is part of CHScanner.
**
** CHScanner is free software: you can redistribute it and/or modify
** it under the terms of the GNU General Public License as published by
** the Free Software Foundation, either version 3 of the License, or
** (at your option) any later version.
**
** CHScanner is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
** GNU General Public License for more details.
**
** You should have received a copy of the GNU General Public License
** along with CHScanner. If not, see <http://www.gnu.org/licenses/>.
*/

#include "StdAfx.h"
#include "MWMIWorker.h"

#include <comdef.h>

//---------------------------------------------------------------------------

MWMIWorker::MWMIWorker(void){
	locator = NULL;
	services = NULL;

	interactiveRun = true;

	CleanUpUserCredentials();

	status = __wwesUnknown;

	coIsInitialized = false;
}

MWMIWorker::~MWMIWorker(void){
	CleanUpUserCredentials();
	DisconnectFromServer();
	CleanUpCom();
}

//---------------------------------------------------------------------------

bool MWMIWorker::SetStatus(HRESULT hResIn)
{
	switch(hResIn){
		case S_OK :
			status = __wwesOK;
			break;

		case E_ACCESSDENIED :
		case WBEM_E_ACCESS_DENIED :
			status = __wwesAccessDenied;
			break;

		case WBEM_E_NOT_FOUND :
			status = __wwesClassDoesNotExists;
			break;

		default :
			status = __wwesOther;
			break;
	}

	return FAILED(hResIn) ? false : true;
}

//---------------------------------------------------------------------------

bool MWMIWorker::GetUserCredentials(void)
{
	CREDUI_INFO cui;
	BOOL saveInfo;
	DWORD res;

	interactiveRun = true;

	CleanUpUserCredentials();

	cui.cbSize = sizeof(CREDUI_INFO);
	cui.hwndParent = NULL;
	cui.pszMessageText = TEXT("Enter credentials (user@domain and password)");
	cui.pszCaptionText = TEXT("Account Informations");
	cui.hbmBanner = NULL;

	saveInfo = FALSE;

	res = CredUIPromptForCredentials(
			&cui,
			TEXT(""),
			NULL,
			0,
			userName,
			CREDUI_MAX_USERNAME_LENGTH + 1,
			userPass,
			CREDUI_MAX_PASSWORD_LENGTH + 1,
			&saveInfo,
			CREDUI_FLAGS_GENERIC_CREDENTIALS | CREDUI_FLAGS_ALWAYS_SHOW_UI | CREDUI_FLAGS_DO_NOT_PERSIST
			);

	SecureZeroMemory(&authIdentity, sizeof(authIdentity));

#ifdef _UNICODE
	authIdentity.User           = (unsigned short*)userName;
	authIdentity.UserLength     = (unsigned long)_tcsnlen(userName, CREDUI_MAX_USERNAME_LENGTH);
	authIdentity.Domain         = NULL;
	authIdentity.DomainLength   = 0;
	authIdentity.Password       = (unsigned short*)userPass;
	authIdentity.PasswordLength = (unsigned long)_tcsnlen(userPass, CREDUI_MAX_PASSWORD_LENGTH);
	authIdentity.Flags          = SEC_WINNT_AUTH_IDENTITY_UNICODE;
#else
	authIdentity.User           = (unsigned char*)userName;
	authIdentity.UserLength     = (unsigned long)_tcsnlen(userName, CREDUI_MAX_USERNAME_LENGTH);
	authIdentity.Domain         = NULL;
	authIdentity.DomainLength   = 0;
	authIdentity.Password       = (unsigned char*)userPass;
	authIdentity.PasswordLength = (unsigned long)_tcsnlen(userPass, CREDUI_MAX_PASSWORD_LENGTH);
	authIdentity.Flags          = SEC_WINNT_AUTH_IDENTITY_ANSI;
#endif

	status = (res==NO_ERROR) ? __wwesOK : __wwesGetCredentials;
	return res==NO_ERROR ? true : false;
}

bool MWMIWorker::GetProcessCredentials(void)
{
	interactiveRun = false;

	CleanUpUserCredentials();
	
	SecureZeroMemory(&authIdentity, sizeof(authIdentity));

#ifdef _UNICODE
	authIdentity.User           = NULL;
	authIdentity.UserLength     = 0;
	authIdentity.Domain         = NULL;
	authIdentity.DomainLength   = 0;
	authIdentity.Password       = NULL;
	authIdentity.PasswordLength = 0;
	authIdentity.Flags          = SEC_WINNT_AUTH_IDENTITY_UNICODE;
#else
	authIdentity.User           = NULL;
	authIdentity.UserLength     = 0;
	authIdentity.Domain         = NULL;
	authIdentity.DomainLength   = 0;
	authIdentity.Password       = NULL;
	authIdentity.PasswordLength = 0;
	authIdentity.Flags          = SEC_WINNT_AUTH_IDENTITY_ANSI;
#endif

	return true;
}

void MWMIWorker::CleanUpUserCredentials(void)
{
	SecureZeroMemory(userName, (CREDUI_MAX_USERNAME_LENGTH + 1) * sizeof(TCHAR));
	SecureZeroMemory(userPass, (CREDUI_MAX_PASSWORD_LENGTH + 1) * sizeof(TCHAR));

	authIdentity.User           = NULL;
	authIdentity.UserLength     = 0;
	authIdentity.Domain         = NULL;
	authIdentity.DomainLength   = 0;
	authIdentity.Password       = NULL;
	authIdentity.PasswordLength = 0;
	authIdentity.Flags          = 0;

	SecureZeroMemory(&authIdentity, sizeof(authIdentity));
}

void MWMIWorker::SetUserCredentials(TCHAR* userNameIn, TCHAR* userPassIn, bool automaticModeIn)
{
	CleanUpUserCredentials();

	interactiveRun = !automaticModeIn;

	memcpy(userName, userNameIn, (CREDUI_MAX_USERNAME_LENGTH + 1) * sizeof(TCHAR));
	memcpy(userPass, userPassIn, (CREDUI_MAX_PASSWORD_LENGTH + 1) * sizeof(TCHAR));

	if(interactiveRun){
		#ifdef _UNICODE
			authIdentity.User           = (unsigned short*)userName;
			authIdentity.UserLength     = (unsigned long)_tcsnlen(userName, CREDUI_MAX_USERNAME_LENGTH);
			authIdentity.Domain         = NULL;
			authIdentity.DomainLength   = 0;
			authIdentity.Password       = (unsigned short*)userPass;
			authIdentity.PasswordLength = (unsigned long)_tcsnlen(userPass, CREDUI_MAX_PASSWORD_LENGTH);
			authIdentity.Flags          = SEC_WINNT_AUTH_IDENTITY_UNICODE;
		#else
			authIdentity.User           = (unsigned char*)userName;
			authIdentity.UserLength     = (unsigned long)_tcsnlen(userName, CREDUI_MAX_USERNAME_LENGTH);
			authIdentity.Domain         = NULL;
			authIdentity.DomainLength   = 0;
			authIdentity.Password       = (unsigned char*)userPass;
			authIdentity.PasswordLength = (unsigned long)_tcsnlen(userPass, CREDUI_MAX_PASSWORD_LENGTH);
			authIdentity.Flags          = SEC_WINNT_AUTH_IDENTITY_ANSI;
		#endif
	}
	else{
		#ifdef _UNICODE
			authIdentity.User           = NULL;
			authIdentity.UserLength     = 0;
			authIdentity.Domain         = NULL;
			authIdentity.DomainLength   = 0;
			authIdentity.Password       = NULL;
			authIdentity.PasswordLength = 0;
			authIdentity.Flags          = SEC_WINNT_AUTH_IDENTITY_UNICODE;
		#else
			authIdentity.User           = NULL;
			authIdentity.UserLength     = 0;
			authIdentity.Domain         = NULL;
			authIdentity.DomainLength   = 0;
			authIdentity.Password       = NULL;
			authIdentity.PasswordLength = 0;
			authIdentity.Flags          = SEC_WINNT_AUTH_IDENTITY_ANSI;
		#endif
	}
}

//---------------------------------------------------------------------------

bool MWMIWorker::InitComAndWmi(void)
{
	HRESULT hRes;

	CleanUpCom();

	hRes = CoInitializeEx(NULL, COINIT_MULTITHREADED);
	if(FAILED(hRes))
		return SetStatus(hRes);
	
	coIsInitialized = true;

	SOLE_AUTHENTICATION_INFO authInfo[2];
	SOLE_AUTHENTICATION_LIST authList;

	SecureZeroMemory(authInfo, 2*sizeof(SOLE_AUTHENTICATION_INFO));

	// NTLM Settings
	authInfo[0].dwAuthnSvc = RPC_C_AUTHN_WINNT;
	authInfo[0].dwAuthzSvc = RPC_C_AUTHZ_NONE;
//	authInfo[0].pAuthInfo  = interactiveRun ? &authIdentity : NULL;
	authInfo[0].pAuthInfo  = &authIdentity;

	// Kerberos Settings
	authInfo[1].dwAuthnSvc = RPC_C_AUTHN_GSS_KERBEROS ;
	authInfo[1].dwAuthzSvc = RPC_C_AUTHZ_NONE;
//	authInfo[1].pAuthInfo  = interactiveRun ? &authIdentity : NULL;
	authInfo[1].pAuthInfo  = &authIdentity;

	authList.cAuthInfo = 2;
	authList.aAuthInfo = authInfo;

	hRes = CoInitializeSecurity( 
		NULL, 
		-1, 
		NULL, 
		NULL, 
		RPC_C_AUTHN_LEVEL_CALL, 
		RPC_C_IMP_LEVEL_IMPERSONATE,
		&authList, 
		EOAC_NONE,
		NULL);

	return SetStatus(hRes);
}

void MWMIWorker::CleanUpCom(void)
{
	DisconnectFromServer();

	if(coIsInitialized){
		CoUninitialize();
		coIsInitialized = false;
	}
}

//---------------------------------------------------------------------------

bool MWMIWorker::ConnectToServer(CString ipAddrIn, CString namespaceIn)
{
	CString str;

	str.Format(_T("\\\\%s%s"), ipAddrIn, namespaceIn);

	return ConnectToServer(str);
}

bool MWMIWorker::ConnectToServer(CString networkResourceIn)
{
	HRESULT hRes;
	BSTR bstrNetworkResource;
	BSTR bstrUser;
	BSTR bstrPass;
	bool err;
	
	hRes = CoCreateInstance(CLSID_WbemLocator, 0, CLSCTX_INPROC_SERVER, IID_IWbemLocator, (LPVOID*)&locator);
	if(FAILED(hRes))
		return SetStatus(hRes);

	err = false;
	try{
		bstrNetworkResource = networkResourceIn.AllocSysString();
	}
	catch(...){
		err = true;
	}
	if(err){
		status = __wwesMemAllocError;
		locator->Release();
		locator = NULL;
		return false;
	}

#ifdef _UNICODE
	bstrUser = SysAllocString(userName);
	bstrPass = SysAllocString(userPass);
#else
	CT2W wUserName(userName);
	CT2W wUserPass(userPass);

	bstrUser = SysAllocString(wUserName);
	bstrPass = SysAllocString(wUserPass);
#endif

	hRes = locator->ConnectServer(
				bstrNetworkResource,
				interactiveRun ? bstrUser : NULL,
				interactiveRun ? bstrPass : NULL,
				0,
				WBEM_FLAG_CONNECT_USE_MAX_WAIT,
				0,
				0,
				&services);
	SysFreeString(bstrUser);
	SysFreeString(bstrPass);

	if(hRes==WBEM_E_LOCAL_CREDENTIALS){
		hRes = locator->ConnectServer(
					bstrNetworkResource,
					NULL,
					NULL,
					0,
					NULL,
					0,
					0,
					&services);
	}
	
	SysFreeString(bstrNetworkResource);

	if(FAILED(hRes)){
		locator->Release();
		locator = NULL;
		return SetStatus(hRes);
	}

	try{
		hRes = CoSetProxyBlanket(
					services,
					RPC_C_AUTHN_WINNT,
					RPC_C_AUTHZ_NONE,
					NULL,
					RPC_C_AUTHN_LEVEL_CALL,
					RPC_C_IMP_LEVEL_IMPERSONATE,
					interactiveRun ? &authIdentity : NULL,
					EOAC_NONE
					);
		/*
		hRes = CoSetProxyBlanket(
					services,
					RPC_C_AUTHN_DEFAULT, //RPC_C_AUTHN_WINNT,
					RPC_C_AUTHZ_DEFAULT, //RPC_C_AUTHZ_NONE,
					COLE_DEFAULT_PRINCIPAL, //NULL,
					RPC_C_AUTHN_LEVEL_DEFAULT, //RPC_C_AUTHN_LEVEL_CALL,
					RPC_C_IMP_LEVEL_IMPERSONATE,
					interactiveRun ? &authIdentity : NULL,
					EOAC_DEFAULT //EOAC_NONE
					);
		*/
	}
	catch(...){
		hRes = -1;
	}

	if(FAILED(hRes))
		DisconnectFromServer();

	return SetStatus(hRes);
}

void MWMIWorker::DisconnectFromServer(void)
{
	if(services!=NULL){
		services->Release();
		services = NULL;
	}
	if(locator!=NULL){
		locator->Release();
		locator = NULL;
	}
}

//---------------------------------------------------------------------------

bool MWMIWorker::ExecuteQuery(CString queryIn)
{
	IEnumWbemClassObject *enumerator;
	BSTR queryStr;
	HRESULT hRes;
	bool err;

	if(services==NULL){
		status = __wwesNotConnected;
		return false;
	}

	enumerator = NULL;

	err = false;
	try{
		queryStr = queryIn.AllocSysString();
	}
	catch(...){
		err = true;
	}
	if(err){
		status = __wwesMemAllocError;
		return false;
	}

	hRes =  services->ExecQuery(
				bstr_t("WQL"),
				queryStr,
				WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
				NULL,
				&enumerator);
	SysFreeString(queryStr);
	if(FAILED(hRes))
		return SetStatus(hRes);

	hRes = CoSetProxyBlanket(
				enumerator,
				RPC_C_AUTHN_WINNT,
				RPC_C_AUTHZ_NONE,
				NULL,
				RPC_C_AUTHN_LEVEL_CALL,
				RPC_C_IMP_LEVEL_IMPERSONATE,
				interactiveRun ? &authIdentity : NULL,
				EOAC_NONE);
	if(FAILED(hRes))
		return SetStatus(hRes);

	IWbemClassObject *classObject;
	ULONG retVal;
	VARIANT vt;
	HRESULT hRes2;

	classObject = NULL;
	retVal = 0;
	while(enumerator){
		hRes = enumerator->Next(WBEM_INFINITE, 1, &classObject, &retVal);
		if(retVal!=0 && classObject!=NULL){
			VariantInit(&vt);

			BSTR propName;
			hRes = classObject->BeginEnumeration(0);
			bool done;
			done = false;
			while(!FAILED(hRes) && !done){
				hRes = classObject->Next(0, &propName, &vt, NULL, NULL);

#ifdef _UNICODE
				OutputDebugString(propName);
#else
				OutputDebugStringW(propName);
#endif

				if(hRes==WBEM_S_NO_MORE_DATA) done = true;
				else{
					OutputDebugString(_T(" : "));

					VariantClear(&vt);
					VariantInit(&vt);
					hRes2 = classObject->Get(propName, 0, &vt, 0, 0);
					if(!FAILED(hRes2)){
						hRes2 = VariantChangeType(&vt, &vt, 0, VT_BSTR);
						if(!FAILED(hRes2)){
#ifdef _UNICODE
							OutputDebugString(vt.bstrVal);
#else
							OutputDebugStringW(vt.bstrVal);
#endif
						}
					}
				}

				OutputDebugString(_T("\n"));

				SysFreeString(propName);
				VariantClear(&vt);
			}
			classObject->EndEnumeration();

//			hRes = classObject->Get(L"Name", 0, &vt, 0, 0);

			VariantClear(&vt);
		}

		if(classObject!=NULL){
			classObject->Release();
			classObject = NULL;
		}

		if(retVal==0){
			if(enumerator!=NULL){
				enumerator->Release();
				enumerator = NULL;
			}
		}
	}

	if(enumerator!=NULL)
		enumerator->Release();

	status = __wwesOK;
	return true;
}

//---------------------------------------------------------------------------

bool MWMIWorker::ExecuteQueryInfo(MWmiInfo* wmiInfo)
{
	IEnumWbemClassObject *enumerator;
	IWbemClassObject *classObject;
	BSTR queryStr;
	HRESULT hRes;
	CString str;
	VARIANT vt;
	int i, cnt;
	unsigned long retCnt;
	bool retVal;
	MWmiInfo *elem;
	MWmiInfo *newElem;
	bool firstElem;
	bool done;
	bool err;

	if(services==NULL){
		status = __wwesNotConnected;
		return false;
	}
	if(wmiInfo==NULL){
		status = __wwesNothingToQuery;
		return false;
	}

	wmiInfo->Initialize();
	str = wmiInfo->GetQueryString();
	if(str.IsEmpty()){
		status = __wwesNothingToQuery;
		return true;
	}

	enumerator = NULL;

	err = false;
	try{
		queryStr = str.AllocSysString();
	}
	catch(...){
		err = true;
	}
	if(err){
		status = __wwesMemAllocError;
		return false;
	}

	hRes =  services->ExecQuery(
				bstr_t("WQL"),
				queryStr,
				WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
				NULL,
				&enumerator);
	SysFreeString(queryStr);
	if(FAILED(hRes))
		return SetStatus(hRes);

	hRes = CoSetProxyBlanket(
				enumerator,
				RPC_C_AUTHN_WINNT,
				RPC_C_AUTHZ_NONE,
				NULL,
				RPC_C_AUTHN_LEVEL_CALL,
				RPC_C_IMP_LEVEL_IMPERSONATE,
				interactiveRun ? &authIdentity : NULL,
				EOAC_NONE);
	if(FAILED(hRes))
		return SetStatus(hRes);

	classObject = NULL;
	retCnt = 0;
	retVal = true;
	firstElem = true;
	done = false;
	while(!done){
		hRes = enumerator->Next(WBEM_INFINITE, 1, &classObject, &retCnt);
		if(SUCCEEDED(hRes) && retCnt==1 && classObject!=NULL){
			if(firstElem){
				firstElem = false;
				elem = wmiInfo;
			}
			else{
				newElem = elem->CreateChild();
				elem->AddChild(newElem);
				elem = newElem;
			}

			if(elem==NULL) retVal = false;
			else{
				cnt = elem->GetPropertyCount();
				for(i=0; i<cnt; i++){
					str = elem->GetPropertyName(i);
					if(!str.IsEmpty()){
						VariantInit(&vt);
						CT2W wStr(str);
						hRes = classObject->Get(wStr, 0, &vt, 0, 0);
						if(SUCCEEDED(hRes))
							retVal &= elem->SetPropertyValue(i, &vt);
						VariantClear(&vt);
					}
				}
			}
		}
		else
			done = true;

		if(classObject!=NULL){
			classObject->Release();
			classObject = NULL;
		}
	}

	if(enumerator!=NULL)
		enumerator->Release();

	status = retVal ? __wwesOK : __wwesOther;
	return retVal;
}

//---------------------------------------------------------------------------
